import os 
import re
import pickle
import cv2

import numpy as np
from tqdm import tqdm


def process_label(alllabelfilepath, xx):
    """
    function: add each action start and end time timestamp(detected by PySlowFast)
    args: alllabelfilepath: the root filepath of each video
          xx: the filename of each video
    return: savetimestamp the named by "timestamp.npy", save each action start and end time timestamp
    e.g.
    alllabelfilepath='./VIDEO_TRANSFER_STATION/'
    alllabelfilepathlist=os.listdir(alllabelfilepath)
    alllabelfilepathlist = sorted(alllabelfilepathlist, key=lambda x: int(re.findall('\d+', x)[0]))

    for i in tqdm(range(len(alllabelfilepathlist))):
        process_label(alllabelfilepath,alllabelfilepathlist[i])
    """
    sketch_label = os.path.join(alllabelfilepath, 
                                xx, 
                                'labels.npy')
    labels = np.load(sketch_label, 
                     allow_pickle=True)
    labels = np.array(labels)

    # Extract detection labels for each frame from labels
    labellist = []
    for i in range(len(labels)):
        if labels[i].numel():
            _, indices = labels[i][0].topk(k=5)
            labellist.append(list(np.array(indices)))

    for i in range(len(labellist)):
        sit = 10
        stand = 11
        walk = 13
        if sit in labellist[i]:
            labellist[i] = 10
        elif stand in labellist[i]:
            labellist[i] = 11
        elif walk in labellist[i]:
            labellist[i] = 13
        else:
            labellist[i] = labellist[i][0]

    for i in range(1, len(labellist) - 1):
        if labellist[i] != labellist[i - 1] and labellist[i] != labellist[i + 1] and labellist[i - 1] == labellist[i + 1]:
            labellist[i] = labellist[i - 1]

    result_index = []
    start = 0
    for i in range(1, len(labellist)):
        if labellist[i] != labellist[i - 1]:
            if i - start > 1:
                result_index.append([start, i])
            else:
                result_index.append([start])
            start = i

    # Process the last continuous segment
    if len(labellist) - start > 1:
        result_index.append([start, len(labellist)])
    else:
        result_index.append([start])

    for i in range(len(result_index)):
        for j in range(len(result_index[i])):
            result_index[i][j] = result_index[i][j] * 64 + 31

    videoframe = os.path.join(alllabelfilepath, xx, 'frame.pkl')

    # Open the pickle file and read the object
    with open(videoframe, 'rb') as f:
        frame = pickle.load(f)

    startandendtimestamp = []
    startandendindex = []
    for i in result_index:
        if len(i) > 1:
            x = []
            xindex = []
            for j in i:
                if j > len(frame) - 1:
                    x.append(frame[len(frame) - 1].timestamp())
                    xindex.append(len(frame) - 1)
                else:
                    x.append(frame[j].timestamp())
                    xindex.append(j)
            startandendindex.append(xindex)
            startandendtimestamp.append(x)

    end2 = frame[-1].timestamp()
    print(end2 == startandendtimestamp[-1][-1])

    if end2 > startandendtimestamp[-1][-1]:
        startandendtimestamp.append([startandendtimestamp[-1][-1], end2])

    startandendtimestamp = np.array(startandendtimestamp)

    savetimestamp = os.path.join(alllabelfilepath, xx, 'timestamp.npy')
    np.save(savetimestamp, startandendtimestamp)



def extract_frames_from_video(SYNCHRO_FOLDER='./SYNCHRONIZATION/'):
    """
    function: extract frames from video

    Args:
        SYNCHRO_FOLDER (str): 包含同步数据的文件夹路径，默认为'./SYNCHRONIZATION/'。

    # Example usage:
    # extract_frames_from_video('./SYNCHRONIZATION/')
    """
    SYNCHRO_FOLDER_LIST = os.listdir(SYNCHRO_FOLDER)
    SYNCHRO_FOLDER_LIST = sorted(SYNCHRO_FOLDER_LIST, key=lambda x: int(re.findall(r'\d+', x)[0]))

    for i in tqdm(range(len(SYNCHRO_FOLDER_LIST))):
        ACTION_TIME = SYNCHRO_FOLDER_LIST[i]

        # 加载时间同步索引
        INIT_SYNC_INDEX = np.load(os.path.join(SYNCHRO_FOLDER, ACTION_TIME, 
                                               'initial_sync_index.npy'), 
                                               allow_pickle=True)

        # 视频文件路径
        VIDEO_FILEPATH = os.path.join(SYNCHRO_FOLDER, 
                                      ACTION_TIME, 
                                      'VIDEO', 
                                      ACTION_TIME, 
                                      'rgb.avi')

        # 图像输出文件夹
        VIDEOTIME_FOLDER = os.path.join('TRANSFER_STATION_VIDEO2IMG', 
                                        ACTION_TIME)
        if not os.path.exists(VIDEOTIME_FOLDER):
            os.mkdir(VIDEOTIME_FOLDER)

            # 打开视频文件
            cap = cv2.VideoCapture(VIDEO_FILEPATH)
            if not cap.isOpened():
                print("Error opening video file")

            fps = cap.get(cv2.CAP_PROP_FPS)
            frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

            for i in range(len(INIT_SYNC_INDEX)):
                ONE_ACTION = INIT_SYNC_INDEX[i]

                VIDEO2IMG_FOLDER = os.path.join(VIDEOTIME_FOLDER, str(i))
                if not os.path.exists(VIDEO2IMG_FOLDER):
                    os.mkdir(VIDEO2IMG_FOLDER)

                for j in range(len(ONE_ACTION)):
                    VIDEO_FPS_INDEX = int(ONE_ACTION[j][0])

                    VIDEO2IMG_OUTPUT_PATH = os.path.join(
                        VIDEO2IMG_FOLDER, str(VIDEO_FPS_INDEX) + '.jpg')

                    if VIDEO_FPS_INDEX < 0 or VIDEO_FPS_INDEX >= frame_count:
                        print("Invalid frame index")
                    else:
                        cap.set(cv2.CAP_PROP_POS_FRAMES, VIDEO_FPS_INDEX)
                        ret, frame = cap.read()

                        if not ret:
                            print("Error reading frame")
                        else:
                            frame_resized = cv2.resize(frame, (250, 250))
                            cv2.imwrite(VIDEO2IMG_OUTPUT_PATH, frame_resized)

            cap.release()


def process_labeled_images(LABELED_FOLDER='./LABELED_IMG/', intev=2):
    """
    处理标记图像文件夹以生成时间戳数据。

    Args:
        LABELED_FOLDER (str): 包含标记图像的文件夹路径，默认为'./LABELED_IMG/'。
        intev (int): 时间间隔（以秒为单位），用于将图像序列划分为时间段，默认为2。
    Returns:
        list: 包含时间戳数据的列表，每个元素包括文件ID、起始帧、结束帧和类别。
    
    Example usage:
    alltime_folderlist = process_labeled_images('./LABELED_IMG/', intev=2)    
    """
    LABELED_FOLDER_LIST = os.listdir(LABELED_FOLDER)
    LABELED_FOLDER_LIST = sorted(LABELED_FOLDER_LIST, key=lambda x: int(re.findall(r'\d+', x)[0]))

    alltime_folderlist = dict()

    for action_imgfoldername in LABELED_FOLDER_LIST:
        action_imgfolderpath = os.path.join(LABELED_FOLDER, action_imgfoldername)

        action_imgfolderpath_list = os.listdir(action_imgfolderpath)
        action_imgfolderpath_list = sorted(action_imgfolderpath_list, key=lambda x: int(re.findall(r'\d+', x)[0]))

        one_action_list = dict()

        for imgfoldername in action_imgfolderpath_list:
            one5sindex = []

            imgfolderpath_list = os.listdir(os.path.join(action_imgfolderpath, imgfoldername))
            imgfolderpath_list = sorted(imgfolderpath_list, key=lambda x: int(re.findall(r'\d+', x)[0]))

            if not len(imgfolderpath_list):
                print(imgfoldername)

            videolen = (int(imgfolderpath_list[-1][:-4]) - int(imgfolderpath_list[0][:-4])) / 30
            intev_time = intev * 30

            if videolen >= intev:
                for j in range(int(videolen // intev)):
                    one5sindex.append([int(imgfolderpath_list[0][:-4]) + (j) * intev_time,
                                       int(imgfolderpath_list[0][:-4]) + (j + 1) * intev_time])
                    one_action_list[imgfoldername] = one5sindex

            else:
                one5sindex.append([int(imgfolderpath_list[0][:-4]), int(imgfolderpath_list[-1][:-4])])
                one_action_list[imgfoldername] = one5sindex

        alltime_folderlist[action_imgfoldername] = one_action_list


    return alltime_folderlist

